%% Disaster model using AMPL
% Code by Sihwan Yang (PhD student in UCLA Economics), 11/05/2020
clear
clc
tic
%% AMPL setting

setting = 1; % David : 1, Sihwan : 2

% Choose elasticity parameters - Create vector of elast.

sigma_s = [0.5, 0.7, 0.95]; % elasticity across consumption
theta1_s = [0.001, 0.1, 0.5]; % elasticity across intermediate
epsilon_s = [0.3, 0.5, 0.7]; % elasticity across VA and intermediate
eta_s = [0.2, 0.5, 0.95]; % elasticity across VA

% very short run : (i,j,u,l) = (1,1,1,1)
% [sigma, theta, epsilon, eta] = [0.5, 0.001, 0.3, 0.2]
% short run : (i,j,u,l) = (2,1,1,1)
% [sigma, theta, epsilon, eta] = [0.7, 0.001, 0.3, 0.2] 
% medium run : (i,j,u,l) = (3,1,3,2)
% [sigma, theta, epsilon, eta] = [0.95, 0.001, 0.7, 0.5]

elas0 = 1; % elas: start
elas1 = 3; % elas: end
theta_sr = [0.7, 0.001, 0.3, 0.2];
theta_mr = [0.95, 0.001, 0.5, 0.5];

if setting == 1
% David
try
run('C:\Program Files\Artelys\AMPL 11.0.20180624\amplapi\examples\matlab\setupOnce.m');
end

else
% Sihwan
addpath('/Users/sihwan/Documents/MATLAB/Baqaee_RA/Darwinian_Frontier_AMPL/amplapi/examples/matlab');
end
%% Load shock
N = 66;

% shock scenario
% 1 -- 10% to 50% universal shock on capital destruction
% 2 -- 10% to 50% shock on manufacturing (8:26) capital
scenario = 1;

if scenario == 1
    Capital_shock = [-0.1, -0.2, -0.3, -0.4, -0.5].*ones(N,1);
elseif scenario == 2
    Capital_shock = [-0.1, -0.2, -0.3, -0.4, -0.5].*[zeros(7,1) ; ones(19,1) ; zeros(40,1)];
end

Np = size(Capital_shock,2);

%% Data: IO matrix(Omega), expenditure share(beta)

% Import IO data matrix
if (exist('IO_data_2018.mat', 'file') ~= 2)
    count = 1;
for year = [(1997:2018) ]
        [~, ~, raw] = xlsread('IOUse_Before_Redefinitions_PRO_1997-2015_Summary+2018.xlsx',num2str(year),'A8:CV94');
        R = cellfun(@(x) ~isnumeric(x) && ~islogical(x),raw); % Find non-numeric cells
        raw(R) = {NaN}; % Replace non-numeric cells
        Data_raw(:,:,count)= reshape([raw{:}],size(raw));
        count = count+1;
end
    [~, ~, indname] = xlsread('IOUse_Before_Redefinitions_PRO_1997-2015_Summary+2018.xlsx','1997','B8:B90');
    indname(cellfun(@(x) ~isempty(x) && isnumeric(x) && isnan(x),indname)) = {''};

    clear ans count R raw year
    save IO_data_2018.mat
else
    load IO_data_2018.mat
end

year = 2015;

IO = Data_raw;
Final = IO(:,99,:);
IO(:,1:2,:) = [];
IO(isnan(IO)) = 0;
IO(IO < 0) = 0;
grossout = squeeze(sum(IO(1:66,1:66,:)))+squeeze(IO(77,1:66,:))+squeeze(IO(79,1:66,:));
labor=(squeeze(IO(77,1:66,:)));
gos=(squeeze(IO(79,1:66,:))); %gross operating surplus

IO(67:end,:,:) = [];
IO(:,67:end,:) = [];

temp = IO(:,:,year-1996);
temp = temp';
Omega = bsxfun(@rdivide, temp, sum(temp,2)); 
K= gos(:,year-1996);
L= labor(:,year-1996);
va_sales=(L+ K);
alphaL=L./va_sales;
alphaK=K./va_sales;
%alphaK = max(alphaK,0); 
N = length(Omega);
beta = Final(1:N,year-1996);
beta(beta<0) = 0; %remove industries with negative implied final sales
beta = beta/sum(beta);
    zero= find(beta==0);

int_sales=squeeze(sum(IO(1:66,1:66,year-1996)))';
int_share=int_sales./(int_sales+va_sales);
va_share=(L+K)./(int_sales+va_sales);
average_mu = 1; 
mu = 1; 

%% Write Omega in standard form - relabeling
%Omega_re= final consumers + N final produc.+ N va. produc. + N int. produc. +  2FACTORS  

labor_temp = labor./grossout;
gos_temp =gos./grossout; %gross operating surplus
alphaL= labor_temp(:,year-1996);
mu = min(mu, 1./alphaL); 
alphaK= gos_temp(:,year-1996) - 1 +1./average_mu';
alphaK = max(alphaK,0); 


F=2*N;
D=1+3*N+F+3; 
% 1 is consumption today, 2:1+N is goods today, N+2:2*N+1 is VA today,
% 2*N+2:3*N+1 is intermediates today, 3*N+2:4*N+1 is labor today, 
% 4*N+2:5*N+1 is capital today, 5*N+1 is HtM consumer, 5*N+2 is the Ricardian
% consumer, 5*N+3 is (aggregate) consumption good tomorrow. 

Omega_re= zeros(D,D);
Omega_re(1,2:N+1)= beta;
Omega_re(2:1+N,1+N+1:2*N+1) = diag(mu.^(-1))*diag(va_share);
Omega_re(2:1+N,1+2*N+1:3*N+1) = diag(mu.^(-1))*diag(int_share);
for x = 1:N
Omega_re(1+N+x, 1+3*N+x)=alphaL(x)/(alphaL(x)+alphaK(x));
Omega_re(1+N+x, 1+4*N+x)=alphaK(x)/(alphaL(x)+alphaK(x));
end
Omega_re(1+2*N+1:1+3*N,2:N+1) = Omega;

Omega_re(1+5*N+1,1)=1;
Omega_re(1+5*N+2,1)=1/2;
Omega_re(1+5*N+2,end)=1/2; 

mu_ones= ones(D,1);
in_mu=mu_ones;
Omega_tilde= diag(in_mu)*Omega_re;
   
Psi_re = inv(eye(D)-Omega_re);
Psi_tilde = inv(eye(D)-Omega_tilde);

%% Evaluate GDP, Welfare

th = 0.1;
mid = 1/th/2+1;
T = length(0:th:1);

Delta_r_GDP_stack = zeros(3,Np,T); % Elasticity/# of shocks
GDP_stack = zeros(3,Np,T);
Hulten_stack = zeros(3,Np,T);
C_postshock_stack = zeros(3,Np,T);

cons_share_stack = zeros(3,Np,D);
lambda_stack = zeros(3,Np,D);
p_stack = zeros(3,Np,D);

for i=1:3 % Elasticity (Short run, Medium run, Cobb-Douglas)
    for j=1:Np % # of shocks

% (1) Choice of Elasticity
if i == 1 % Short run
    sigma = theta_sr(1);
    theta1 = theta_sr(2)*ones(N,1);
    epsilon = theta_sr(3)*ones(N,1);
    eta = theta_sr(4)*ones(N,1);
elseif i == 2 % Medium run
    sigma = theta_mr(1);
    theta1 = theta_mr(2)*ones(N,1);
    epsilon = theta_mr(3)*ones(N,1);
    eta = theta_mr(4)*ones(N,1);
elseif i == 3 % Cobb-Douglas
    sigma = 0.99; 
    theta1 = 0.99*ones(N,1); 
    epsilon = 0.99*ones(N,1);
    eta = 0.99*ones(N,1); 
end

factor_ela=.5*ones(F,1);
rho = .99;
theta = cat(1,sigma,epsilon,eta, theta1,factor_ela,.9,rho,.9);

% (2) Choice of shocks
if Np == 1
    shock_supply = -Capital_shock(:,j);
else
    if j==1
        shock_supply = -Capital_shock(:,j);
    else
        shock_supply = -Capital_shock(:,j);
        shock_supply_pre = -Capital_shock(:,j-1);
    end
end
    
% Create a categorical variable to set if conditions in constraints in AMPL
factor = ones(D,1); % one for goods
factor(1+3*N+1:1+5*N, 1) = 0; % zero for factors
factor(end) = 0; % Since consumption tomorrow is treated like a factor. 
factor(end-2) = 3; % since the second and third to last rows are consumers. 
factor(end-1) = 2; % since the second and third to last rows are consumers. 

keynes = ones(D,1);  % one for goods, zero for factors
keynes(1+3*N+1:1+4*N, 1) = 0; % zero for flexible factors: labor
keynes(1+4*N+1:1+5*N, 1) = 0; % zero for flexible factors: capital
keynes(1+5*N+3, 1) = 0; % zero for tomorrow's consumption

% ownership of the factors
chi = zeros(D,1);

% Determin price elasticity of factor supply 
phi = zeros(D,1); 

% Change Ind format for set
Ind= (1:D)';
Ind = string(Ind);
Ind = cellstr(Ind);
Indu = cellstr(Ind);

init_lambda = Psi_re(1,:)';
init_lambda(1) = 1;
init_lambda(end-2) = (Psi_re(1,:)*chi)*init_lambda(1);
init_lambda(end-1) = 2*(1-(Psi_re(1,:)*chi))*init_lambda(1);
init_lambda(end) = 0.5*init_lambda(end-1); 

Index = zeros(1,D);
foo = sort(Ind);
for v = 1:D
    Index(1,v) = str2num(foo{v});
end
foo = sortrows([Index' (1:D)']);
Index = foo(:,2);


shock_step = 0;
Delta_r_GDP = zeros(T,1);
GDP = zeros(T,1);
Hulten = zeros(T,1);
C_postshock = zeros(T,1);

nominal_GDP_s = [];
c_shares_s = [];
prices_s = [];

if Np == 1
    init_lambda1 = init_lambda; 
    init_p = ones(D,1);
else
    if j==1
        init_lambda1 = init_lambda;
        init_p = ones(D,1);
    else
        init_lambda1 = init_lambda1_temp;
        init_p = init_p_temp;
    end
end
        for t=0:th:1 % Shock Intensity (0, 0.5, 1)

shock_step = shock_step + 1;  

if setting == 1
% David
     try
ampl.close()
     end
ampl = AMPL('C:\Program Files\Artelys\AMPL 11.0.20180624');
ampl.setOption('solver', 'knitroampl')

else
     try
ampl.close()
end
setupOnce
ampl = AMPL('/Users/sihwan/Documents/MATLAB/AMPL_Mac');
ampl.setOption('solver', 'knitro') 
end

% Choose A and mu
% (3) Choice of Shock Order/Intensity

A = ones(D,1);
B = ones(D,D);

if Np == 1
    A(1+4*N+1:1+5*N) =(1-t*shock_supply);
else
    if j==1
        A(1+4*N+1:1+5*N) =(1-t*shock_supply);
    else
        A(1+4*N+1:1+5*N) =(1-t*shock_supply-(1-t)*shock_supply_pre);
    end
end
in_mu=ones(D,1);
mu=ones(D,1);

% Select model
ampl.read ('standard_form_neoclassical.mod')

% Send data to AMPL
df = DataFrame(2, 'Ind', 'Indu', 'Omega');
df.setMatrix(Omega_tilde, Ind, Indu)
ampl.eval('set S1; set S2; param TheParameter{S1, S2};');
ampl.getSet('S1').setValues(Ind);
ampl.getSet('Indu').setValues(Indu);
ampl.setData(df);       

df = DataFrame(2, 'Ind', 'Indu', 'B');
df.setMatrix(B, Ind, Indu)
ampl.eval('set S11; set S12; param TheParameter1{S11, S12};');
ampl.getSet('S11').setValues(Ind);
ampl.getSet('Indu').setValues(Indu);
ampl.setData(df);    

 df1 = DataFrame(1, 'Ind','A', 'mu', 'theta','factor', 'in_mu','init_p','init_lambda','init_lambda1','keynes','chi','phi');
    df1.setColumn('Ind', Ind );
    df1.setColumn('A', A );
    df1.setColumn('mu', mu );
    df1.setColumn('theta', theta ); 
    df1.setColumn('factor', factor);
    df1.setColumn('in_mu', in_mu);
    df1.setColumn('init_p', init_p);
    df1.setColumn('init_lambda', init_lambda);
    df1.setColumn('init_lambda1', init_lambda1);
    df1.setColumn('keynes', keynes);
    df1.setColumn('chi',chi);
    df1.setColumn('phi',phi);
  ampl.setData(df1, 'Ind');


% Solve  
  ampl.solve();
  
% Get solutions from AMPL

ampl.eval('var return_code:=solve_result_num;')
ampl.eval('display return_code;')
return_code = ampl.getVariable('return_code');
df = return_code.getValues('val');
return_code = df.getColumnAsDoubles('return_code.val');

p = ampl.getVariable('p');
df = p.getValues('val');
p = df.getColumnAsDoubles('p.val'); 

lambda = ampl.getVariable('lambda');
df = lambda.getValues('val');
lambda = df.getColumnAsDoubles('lambda.val'); 
 
cons_share = ampl.getVariable('cons_share');
df = cons_share.getValues('val');
cons_share = df.getColumnAsDoubles('cons_share.val');

cons_share = cons_share(Index);
init_p = p(Index);
init_lambda1 = lambda(Index); 

 if return_code == 0 || return_code == 100
     nominal_GDP_s = [nominal_GDP_s, lambda(1)];
     c_shares_s = [c_shares_s, cons_share(2:67)];
     prices_s = [prices_s, init_p(2:67)];
     c_init = init_lambda(1).*Omega_re(1,:)';
     c = lambda(1)/p(1)*(B(1,:).^sigma)'.*Omega_re(1,:)'.*(p(Index)./p(1)).^(-sigma);
     
     if size(c_shares_s,2) == 1
         Delta_r_GDP(shock_step,1) = 0;
     else
        Delta_r_GDP(shock_step,1) = 1-exp(sum(diff(log(nominal_GDP_s))-sum((c_shares_s(:,1:end-1)+c_shares_s(:,2:end))/2.*(diff(log(prices_s)')'))));
     end
     GDP(shock_step,1) = lambda(1)/p(1);
     Hulten(shock_step,1) = exp(Psi_re(1,:)*log(A));
     C_postshock(shock_step,1) = (sum(((B(1,Omega_re(1,:)>0))'.*Omega_re(1,(Omega_re(1,:)>0))'.*(c(Omega_re(1,:)>0)./c_init(Omega_re(1,:)>0)).^((sigma-1)/(sigma))))^((sigma)/(sigma-1)));     
 else
     Delta_r_GDP(shock_step,1) = NaN;
     GDP(shock_step,1) = NaN;
     Hulten(shock_step,1) = NaN;
     C_postshock(shock_step,1) = NaN;
 end

        end
        
        Delta_r_GDP_stack(i,j,:) = Delta_r_GDP';
        GDP_stack(i,j,:) = GDP';
        Hulten_stack(i,j,:) = Hulten';
        C_postshock_stack(i,j,:) = C_postshock';
        cons_share_stack(i,j,:) = cons_share;
        lambda_stack(i,j,:) = init_lambda1;
        p_stack(i,j,:) = init_p;
        
        init_lambda1_temp = init_lambda1;
        init_p_temp = init_p;
        
    end
end
ampl.close();

%% Store result
filename = 'Simulation_replicate.xlsx';

writematrix(reshape(permute(Delta_r_GDP_stack,[3 1 2]),T,3*Np),filename,'Sheet','Delta_r_GDP','Range','C8');
writematrix(reshape(permute(GDP_stack,[3 1 2]),T,3*Np),filename,'Sheet','Welfare','Range','C8');
writematrix(reshape(permute(Hulten_stack,[3 1 2]),T,3*Np),filename,'Sheet','Hulten','Range','C8');
writematrix(reshape(permute(C_postshock_stack,[3 1 2]),T,3*Np),filename,'Sheet','C_postshock','Range','C8');
writematrix([init_lambda(2:N+1), reshape(permute(lambda_stack(:,:,2:N+1),[3 1 2]),N,3*Np)],filename,'Sheet','lambda','Range','C8');
writematrix([init_p(2:N+1), reshape(permute(p_stack(:,:,2:N+1),[3 1 2]),N,3*Np)],filename,'Sheet','price','Range','C8');

rGDP_pt = zeros(Np*(T-1)+1,3);
Welfare_pt = zeros(Np*(T-1)+1,3);
Hulten_pt = zeros(Np*(T-1)+1,3);
C_postshock_pt = zeros(Np*(T-1)+1,3);

for i=1:3
    temp_rGDP = squeeze(1-Delta_r_GDP_stack(i,:,:))';
    temp_Welfare = squeeze(GDP_stack(i,:,:))';
    temp_Hulten =  squeeze(Hulten_stack(i,:,:))';
    temp_C_postshock =  squeeze(C_postshock_stack(i,:,:))';
    
    for j=1:Np
        if j == 1
            rGDP_pt(1:T,i) = temp_rGDP(:,j);
            Welfare_pt(1:T,i) = temp_Welfare(:,j);
            Hulten_pt(1:T,i) = temp_Hulten(:,j);
            C_postshock_pt(1:T,i) = temp_C_postshock(:,j);
            rGDP_last = rGDP_pt(T,i);
        else
            rGDP_pt((T-1)*(j-1)+2:(T-1)*j+1,i) = rGDP_last*temp_rGDP(2:T,j);
            Welfare_pt((T-1)*(j-1)+2:(T-1)*j+1,i) = temp_Welfare(2:T,j);
            Hulten_pt((T-1)*(j-1)+2:(T-1)*j+1,i) = temp_Hulten(2:T,j);
            C_postshock_pt((T-1)*(j-1)+2:(T-1)*j+1,i) = temp_C_postshock(2:T,j);
            rGDP_last = rGDP_pt((T-1)*j+1,i);
        end
    end
end

if scenario == 1 % Universal shock on capital destruction

writematrix(rGDP_pt,filename,'Sheet','Universal','Range','C8');
writematrix(Welfare_pt,filename,'Sheet','Universal','Range','F8');
writematrix(Hulten_pt,filename,'Sheet','Universal','Range','I8');
writematrix(C_postshock_pt,filename,'Sheet','Universal','Range','L8');

grid = (0:1:Np*(T-1))';
figure(1)
plot(grid,rGDP_pt(:,1),'k','Linewidth',2)
hold
plot(grid,rGDP_pt(:,2),'r','Linewidth',2)
plot(grid,rGDP_pt(:,3),'g','Linewidth',2)
plot(grid,Hulten_pt(:,1),'b','Linewidth',2)
hold
title('Universal Capital Destruction')
legend('Short-run','Medium-run','Cobb-Douglas','Hulten')

elseif scenario == 2 % shock on manufacturing

% writematrix(rGDP_pt,filename,'Sheet','Manufacturing','Range','C8');
% writematrix(Welfare_pt,filename,'Sheet','Manufacturing','Range','F8');
% writematrix(Hulten_pt,filename,'Sheet','Manufacturing','Range','I8');
% writematrix(C_postshock_pt,filename,'Sheet','Manufacturing','Range','L8');

grid = (0:1:Np*(T-1))';
figure(1)
plot(grid,rGDP_pt(:,1),'k','Linewidth',2)
hold
plot(grid,rGDP_pt(:,2),'r','Linewidth',2)
plot(grid,rGDP_pt(:,3),'g','Linewidth',2)
plot(grid,Hulten_pt(:,1),'b','Linewidth',2)
hold
title('Manufacturing-Only Capital Destruction')
legend('Short-run','Medium-run','Cobb-Douglas','Hulten')

end

lambda_heat = zeros(N,Np,3);
p_heat = zeros(N,Np,3);
capital_heat = zeros(N,Np,3);
r_heat = zeros(N,Np,3);

for i=1:3
    lambda_heat_init = init_lambda(2:N+1);
    p_heat_init = ones(N,1);
    capital_heat_init = init_lambda(4*N+2:5*N+1);
    r_heat_init = ones(N,1);
    for j=1:Np
        lambda_heat(:,j,i) = (squeeze(lambda_stack(i,j,2:N+1))./lambda_heat_init);
        p_heat(:,j,i) = (squeeze(p_stack(i,j,2:N+1))./p_heat_init);
        capital_heat(:,j,i) = (squeeze(lambda_stack(i,j,4*N+2:5*N+1))./capital_heat_init);
        r_heat(:,j,i) = (squeeze(p_stack(i,j,4*N+2:5*N+1))./r_heat_init);
    end
end

figure(2)
subplot(1,3,1)
heatmap(lambda_heat(:,:,1),'Colormap',hot);
title('Sales share change - Short-run')
subplot(1,3,2)
heatmap(lambda_heat(:,:,2),'Colormap',hot);
title('Sales share change - Medium-run')
subplot(1,3,3)
heatmap(lambda_heat(:,:,3),'Colormap',hot);
title('Sales share change - Cobb-Douglas')

figure(3)
subplot(1,3,1)
heatmap(p_heat(:,:,1),'Colormap',hot);
title('Price change - Short-run')
subplot(1,3,2)
heatmap(p_heat(:,:,2),'Colormap',hot);
title('Price change - Medium-run')
subplot(1,3,3)
heatmap(p_heat(:,:,3),'Colormap',hot);
title('Price change - Cobb-Douglas')

figure(4)
subplot(1,3,1)
heatmap(lambda_heat(:,:,1)./p_heat(:,:,1),'Colormap',hot);
title('Quantity change - Short-run')
subplot(1,3,2)
heatmap(lambda_heat(:,:,2)./p_heat(:,:,2),'Colormap',hot);
title('Quantity change - Medium-run')
subplot(1,3,3)
heatmap(lambda_heat(:,:,3)./p_heat(:,:,3),'Colormap',hot);
title('Quantity change - Cobb-Douglas')

figure(5)
subplot(1,3,1)
heatmap(capital_heat(:,:,1),'Colormap',hot);
title('Capital share change - Short-run')
subplot(1,3,2)
heatmap(capital_heat(:,:,2),'Colormap',hot);
title('Capital share change - Medium-run')
subplot(1,3,3)
heatmap(capital_heat(:,:,3),'Colormap',hot);
title('Capital share change - Cobb-Douglas')

figure(6)
subplot(1,3,1)
heatmap(r_heat(:,:,1),'Colormap',hot);
title('Capital price change - Short-run')
subplot(1,3,2)
heatmap(r_heat(:,:,2),'Colormap',hot);
title('Capital price change - Medium-run')
subplot(1,3,3)
heatmap(r_heat(:,:,3),'Colormap',hot);
title('Capital price change - Cobb-Douglas')

figure(7)
subplot(1,3,1)
heatmap(capital_heat(:,:,1)./r_heat(:,:,1),'Colormap',hot);
title('Capital Quantity change - Short-run')
subplot(1,3,2)
heatmap(capital_heat(:,:,2)./r_heat(:,:,2),'Colormap',hot);
title('Capital Quantity change - Medium-run')
subplot(1,3,3)
heatmap(capital_heat(:,:,3)./r_heat(:,:,3),'Colormap',hot);
title('Capital Quantity change - Cobb-Douglas')

if scenario == 1
writematrix(lambda_heat(:,:,1),filename,'Sheet','Universal','Range','Q8');
writematrix(lambda_heat(:,:,2),filename,'Sheet','Universal','Range','W8');
writematrix(lambda_heat(:,:,3),filename,'Sheet','Universal','Range','AC8');

elseif scenario == 2
% writematrix(lambda_heat(:,:,1),filename,'Sheet','Manufacturing','Range','Q8');
% writematrix(lambda_heat(:,:,2),filename,'Sheet','Manufacturing','Range','W8');
% writematrix(lambda_heat(:,:,3),filename,'Sheet','Manufacturing','Range','AC8');
    
end
